import os
import yaml
import xarray as xr
from config import conf
# from del_file import finish
from del_file import delete_dir, delete_nc

try:
    from log_own import logger
except ImportError:
    import sys
    sys.path.append(".")
    sys.path.append("../")
    sys.path.append("../../")
    from log_own import logger


def save_var_to_yaml(file):
    var_list = []
    var_info = {}

    if os.path.basename(file).split(".")[0].endswith("-processed"):
        data_type = os.path.basename(file).split(".")[0][18:25]
    else:
        data_type = os.path.basename(file).split(".")[0][-7:]  # 获取文件名中的数据类型，如"wave-fc"

    ds = xr.open_dataset(file)
    for key in ds.keys():
        var_list.append(key)
        print(key)

    var_info[f"{data_type}"] = var_list
    print(var_info)
    with open("config.yml", "a") as f:
        yaml.dump(var_info, f, default_flow_style=True)


def get_var_list(file):

    # 根据文件名获取配置文件属性值，加入列表var_list中
    if os.path.basename(file).split(".")[0].endswith("-processed"):
        # data_type = os.path.basename(file).split(".")[0][18:25]  # "enfo-ef"
        data_type = os.path.basename(file).split(".")[0][-17:-10]  # "enfo-ef"
    else:
        data_type = os.path.basename(file).split(".")[0][-7:]  # 获取文件名中的数据类型，如"wave-fc"

    var_list = conf[f"{data_type}"]
    return var_list


def produce_nc(file):
    """
    把 grib文件的变量依次抽取生成 nc文件，存放在以变量名命名的目录中
    @param file: grib文件绝对路径
    """
    if file.endswith("enfo-ef.grib2"):
        return
    try:
        var_list = get_var_list(file)
    except KeyError as e:
        logger.error(e)
        logger.warning("跳过 " + file)
        return

    # 遍历指定变量名列表var_list，每个变量单独生成一个nc文件
    for var in var_list:
        if var == "10uv" or var == "uv":
            produce_wind(var, file)
        elif var == "wave":
            produce_wave(var, file)
        else:
            ds = xr.open_dataset(file, filter_by_keys={'shortName': f"{var}"})
            if var == "2t":
                var = "t2m"
            print(ds.variables)
            data = ds[f"{var}"]
            new_ds = xr.Dataset({f"{var}": data})

            # 根据变量名和起报时间生成数据存储目录
            forecast_time = os.path.basename(file)[:14]
            store_dir = os.path.join(conf["store_nc_dir"], var, forecast_time)  # C:\Users\20180017\Desktop\ncTest\output + \swh
            if not os.path.exists(store_dir):
                os.makedirs(store_dir)

            file_name = os.path.basename(file).split(".")[0] + f"-{var}" + ".nc"  # nc文件名
            file_path = os.path.join(store_dir, file_name)  # nc文件绝对路径  C:\Users\20180017\Desktop\ncTest\output\swh +
            new_ds.to_netcdf(file_path)
            logger.info("已生成：" + file_path)


def produce_wind(var, file):
    if var == "10uv":
        ds_u = xr.open_dataset(file, filter_by_keys={'shortName': "10u"})
        ds_v = xr.open_dataset(file, filter_by_keys={'shortName': "10v"})
        u_data = ds_u["u10"]
        v_data = ds_v["v10"]
        new_ds = xr.Dataset({"10u": u_data, "10v": v_data})
    elif var == "uv":
        ds_u = xr.open_dataset(file, filter_by_keys={'shortName': "u"})
        ds_v = xr.open_dataset(file, filter_by_keys={'shortName': "v"})
        u_data = ds_u["u"]
        v_data = ds_v["v"]
        new_ds = xr.Dataset({"u": u_data, "v": v_data})

    # 根据变量名和起报时间生成数据存储目录
    forecast_time = os.path.basename(file)[:14]
    store_dir = os.path.join(conf["store_nc_dir"], var, forecast_time)
    if not os.path.exists(store_dir):
        os.makedirs(store_dir)

    file_name = os.path.basename(file).split(".")[0] + f"-{var}" + ".nc"
    file_path = os.path.join(store_dir, file_name)
    new_ds.to_netcdf(file_path)
    logger.info("已生成：" + file_path)


def produce_wave(var, file):
    ds_mwp = xr.open_dataset(file, filter_by_keys={'shortName': "mwp", 'dataType': 'cf'})
    ds_mwd = xr.open_dataset(file, filter_by_keys={'shortName': "mwd", 'dataType': 'cf'})
    ds_swh = xr.open_dataset(file, filter_by_keys={'shortName': "swh", 'dataType': 'cf'})
    mwp_data = ds_mwp["mwp"]
    mwd_data = ds_mwd["mwd"]
    swh_data = ds_swh["swh"]
    new_ds = xr.Dataset({"mwp": mwp_data, "mwd": mwd_data, "swh": swh_data})

    # 根据变量名和起报时间生成数据存储目录
    forecast_time = os.path.basename(file)[:14]
    store_dir = os.path.join(conf["store_nc_dir"], var, forecast_time)
    if not os.path.exists(store_dir):
        os.makedirs(store_dir)

    file_name = os.path.basename(file).split(".")[0] + "-wave" + ".nc"
    file_path = os.path.join(store_dir, file_name)
    new_ds.to_netcdf(file_path)
    logger.info("已生成：" + file_path)


def do_produce():
    # 等待delete函数执行完毕再开始执行
    # while True:
    #     if finish:
    #         break

    # delete_dir()
    # delete_nc()
    target = ["enfo", "oper", "waef", "wave", "scda", "scwv"]
    store_grib_dir = conf["store_grib_dir"]
    for root, dirs, files in os.walk(store_grib_dir):
        for dir in dirs:
            if dir in target:
                data_dir = os.path.join(root, dir)  # ......\oper
                try:
                    # produce_nc(file_path)
                    print(data_dir)
                except KeyError as e:
                    logger.error(e)
                    logger.error("转换失败")


if __name__ == '__main__':
    pass
    # save_var_to_yaml("data/20230624060000-0h-scda-fc.grib2")
    # produce_nc(r"C:\Users\20180017\Desktop\ncTest\20230801120000-51h-oper-fc.grib2")
    # produce_nc("data/20230624000000-360h-waef-ef.grib2")
    do_produce()
